package com.wgh.demo.ai.service.impl;

import com.wgh.demo.ai.config.AiModelConfig;
import com.wgh.demo.ai.enums.AiModelType;
import com.wgh.demo.ai.service.AiModelService;
import com.google.gson.Gson;
import okhttp3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

@Service
public class DeepSeekAiModelServiceImpl implements AiModelService {
    private static final Logger logger = LoggerFactory.getLogger(DeepSeekAiModelServiceImpl.class);
    private static final MediaType JSON = MediaType.get("application/json; charset=utf-8");

    @Autowired
    private AiModelConfig aiModelConfig;

    // 创建带有超时设置的OkHttpClient
    private final OkHttpClient client = new OkHttpClient.Builder()
            .connectTimeout(30, java.util.concurrent.TimeUnit.SECONDS)  // 连接超时
            .readTimeout(60, java.util.concurrent.TimeUnit.SECONDS)     // 读取超时
            .writeTimeout(30, java.util.concurrent.TimeUnit.SECONDS)    // 写入超时
            .build();
    private final Gson gson = new Gson();

    @Override
    public Mono<String> generateText(String prompt) {
        return Mono.fromCallable(() -> {
            // 构建请求参数
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", "deepseek-chat"); // 使用DeepSeek对话模型
            requestBody.put("temperature", 0.7);
            requestBody.put("max_tokens", 1000);

            List<Map<String, String>> messages = new ArrayList<>();
            Map<String, String> message = new HashMap<>();
            message.put("role", "user");
            message.put("content", prompt);
            messages.add(message);

            requestBody.put("messages", messages);

            // 构建请求
            Request request = new Request.Builder()
                    .url(aiModelConfig.getDeepseek().getEndpoint())
                    .post(RequestBody.create(gson.toJson(requestBody), JSON))
                    .addHeader("Authorization", "Bearer " + aiModelConfig.getDeepseek().getApiKey())
                    .build();

            // 发送请求
            try (Response response = client.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    logger.error("DeepSeek大模型请求失败，状态码: {}", response.code());
                    throw new RuntimeException("DeepSeek大模型请求失败，状态码: " + response.code());
                }

                // 解析响应
                String responseBody = response.body() != null ? response.body().string() : null;
                if (responseBody == null) {
                    logger.error("DeepSeek大模型响应为空");
                    throw new RuntimeException("DeepSeek大模型响应为空，请稍后再试");
                }

                Map<String, Object> responseMap = gson.fromJson(responseBody, Map.class);
                if (responseMap.containsKey("error")) {
                    Map<String, Object> errorMap = (Map<String, Object>) responseMap.get("error");
                    String errorMessage = errorMap.get("message").toString();
                    logger.error("DeepSeek大模型响应错误: {}", errorMessage);
                    throw new RuntimeException("DeepSeek大模型响应错误: " + errorMessage);
                }

                if (responseMap.containsKey("choices")) {
                    List<Map<String, Object>> choices = (List<Map<String, Object>>) responseMap.get("choices");
                    if (!choices.isEmpty()) {
                        Map<String, Object> choice = choices.get(0);
                        Map<String, Object> messageMap = (Map<String, Object>) choice.get("message");
                        return messageMap.get("content").toString();
                    }
                }

                logger.error("DeepSeek大模型响应格式异常: {}", responseBody);
                throw new RuntimeException("DeepSeek大模型响应格式异常，请稍后再试");
            }
        }).onErrorResume(e -> {
            logger.error("调用DeepSeek大模型失败: {}", e.getMessage(), e);
            return Mono.just("调用DeepSeek大模型失败: " + e.getMessage());
        });
    }

    @Override
    public String getModelName() {
        return AiModelType.DEEPSEEK.getCode();
    }
}